"""
Langgraph 处理任务流
"""
from typing import TypedDict, List, Dict, Any, Annotated

from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.prompts import ChatPromptTemplate
from pydantic import BaseModel, Field

from model import get_model

from langgraph.graph import StateGraph  # 整个Graph的图
from langgraph.graph import END  # 结束的节点
from langgraph.prebuilt import ToolNode  # 工具节点
from langgraph.checkpoint.memory import MemorySaver  # 保存节点之间所有消息,类似于history (回退)

# --------------- 绑定tools --------------
model = get_model()


# 以下的功能提示词，可以参考智谱文档描述：https://www.bigmodel.cn/dev/howuse/prompt
# --------------- 制定多计划的节点 --------------
class Plans(BaseModel):
    plans: List[str] = Field("如果没有任务可划分,该内容为空")


take_plans_prompt = ChatPromptTemplate.from_messages(
    [
        (
            # 定义系统消息，为生成计划提供规则和要求
            "system",
            """对于给定的目标，提出一个简单的逐步计划。这个计划应该包含独立的任务，
            如果正确执行将得出正确的答案。不要添加任何多余的步骤。最后一步的结果应
            该是最终答案。确保每一步都有所有必要的信息 - 不要跳过步骤。""",
        ),
        # 定义占位符消息，用于后续填充实际的用户输入
        ("user", "{messages}"),
    ]
)


class GraphState(TypedDict):
    title: str
    plans: List[str]
    content: List[str]
    finish: List[str]


def take_many_plans(state: GraphState):
    title = state['title']
    plans_chain = take_plans_prompt | model.with_structured_output(schema=Plans)
    response = plans_chain.invoke({'messages': title})
    print("=" * 20, "Agent1", "=" * 20)
    plans = response.__dict__['plans']
    print(plans)
    return {
        "title": title,
        "plans": plans,
        "content": [],
        "finish": []
    }


#  --------------- 重新制定任务节点 --------------

replanner_prompt = ChatPromptTemplate.from_template(
    """
        对于给定的目标，提出一个简单的逐步计划，步骤必须越少越好。这个计划应该包含独立的任务，如果正确执行将得出正确的答案。不要添加任何多余的步骤。最后一步的结果应该是最终答案。确保每一步都有所有必要的信息 - 不要跳过步骤。
        你的主题是是：
        {title}
        你的原计划是：
        {plans}
        你目前已完成的步骤是：
        {finish}
        返回剩下的步骤，不要返回无原计划不相关内容，如果没有剩余步骤，则返回“END”
    """
)


# 相应地更新你的计划。如果不需要更多步骤并且可以返回给用户，那么就这样回应。如果需要，填写计划。只添加仍然需要完成的步骤。不要返回已完成的步骤作为计划的一部分。


def re_take_many_plans(state: GraphState):
    print("=" * 20, "Agents3收到的信息", "=" * 20)
    title = state['title']
    plans = state['plans']
    content = state['content']
    finish = state['finish']
    print("==当前目标", title)
    print("==任务信息", plans)
    print("==已完成的信息", finish)
    plans_chain = replanner_prompt | model.with_structured_output(schema=Plans)
    response = plans_chain.invoke({"title": title, "plans": plans, "finish": finish})
    print("=" * 20, "Agent3", "=" * 20)
    plans = response.__dict__['plans']
    print(plans)
    return {
        "title": title,
        "plans": plans,
        "content": content,
        "finish": finish
    }


def should_end(state: GraphState):
    plans = state['plans']
    if len(plans) == 0 or 'END' in plans:
        return END
    else:
        return "agent2"


# --------------- 制定单任务执行的节点 --------------


def take_one_task(state: GraphState):
    print("=" * 20, "Agent2", "=" * 20)
    # print(state['messages'])
    title = state['title']
    content = state['content']
    plans = state['plans']
    finish = state['finish']
    first_plan = plans[0]
    response = model.invoke([HumanMessage(first_plan)])
    content.append(response.content)
    finish.append(first_plan)
    return {
        "title": title,
        "plans": plans,
        "content": content,
        "finish": finish
    }


workflow = StateGraph(GraphState)
workflow.add_node("agent1", take_many_plans)
workflow.add_node("agent2", take_one_task)
workflow.add_node("agent3", re_take_many_plans)
workflow.set_entry_point("agent1")  # 设置起点位置

workflow.add_edge("agent1", "agent2")
workflow.add_edge("agent2", "agent3")
workflow.add_conditional_edges(
    "agent3",
    should_end
)

app = workflow.compile(checkpointer=MemorySaver())  # 构建Langgraph应用

final_state = app.invoke(
    {"title": [HumanMessage("把大象关进冰箱")]},
    config={"configurable": {"thread_id": 42}}
)
print(final_state)
